{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "# 利用 Estimator 进行多工作器训练\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://tensorflow.google.cn/tutorials/distribute/multi_worker_with_estimator\"><img src=\"https://tensorflow.google.cn/images/tf_logo_32px.png\" />在 tensorflow.google.cn 上查看</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/tutorials/distribute/multi_worker_with_estimator.ipynb\"><img src=\"https://tensorflow.google.cn/images/colab_logo_32px.png\" />在 Google Colab 运行</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/tutorials/distribute/multi_worker_with_estimator.ipynb\"><img src=\"https://tensorflow.google.cn/images/GitHub-Mark-32px.png\" />在 Github 上查看源代码</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/tutorials/distribute/multi_worker_with_estimator.ipynb\"><img src=\"https://tensorflow.google.cn/images/download_logo_32px.png\" />下载此 notebook</a>\n",
        "  </td>\n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GEe3i16tQPjo"
      },
      "source": [
        "Note: 我们的 TensorFlow 社区翻译了这些文档。因为社区翻译是尽力而为， 所以无法保证它们是最准确的，并且反映了最新的\n",
        "[官方英文文档](https://tensorflow.google.cn/?hl=en)。如果您有改进此翻译的建议， 请提交 pull request 到\n",
        "[tensorflow/docs](https://github.com/tensorflow/docs) GitHub 仓库。要志愿地撰写或者审核译文，请加入\n",
        "[docs-zh-cn@tensorflow.org Google Group](https://groups.google.com/a/tensorflow.org/forum/#!forum/docs-zh-cn)。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xHxb-dlhMIzW"
      },
      "source": [
        "## 概述\n",
        "本教程展示了在训练分布式多工作器（worker）时，如何使用 `tf.distribute.Strategy`。如果你的代码使用了 `tf.estimator`，而且你也对拓展单机以获取高性能有兴趣，那么这个教程就是为你准备的。\n",
        "\n",
        "在开始之前，请先阅读 [`tf.distribute.Strategy` 指南](../../guide/distribute_strategy.ipynb)。同样相关的还有 [使用多 GPU 训练教程](./keras.ipynb)，因为在这个教程里也使用了相同的模型。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MUXex9ctTuDB"
      },
      "source": [
        "## 创建\n",
        "\n",
        "首先，设置好 TensorFlow 以及将会用到的输入模块。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bnYxvfLD-LW-"
      },
      "outputs": [],
      "source": [
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "tfds.disable_progress_bar()\n",
        "\n",
        "import os, json"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hPBuZUNSZmrQ"
      },
      "source": [
        "## 输入函数\n",
        "\n",
        "本教程里我们使用的是 [TensorFlow 数据集（TensorFlow Datasets）](https://tensorflow.google.cn/datasets)里的 MNIST 数据集。本教程里的代码和 [使用多 GPU 训练教程](./keras.ipynb) 类似，但有一个主要区别：当我们使用 Estimator 进行多工作器训练时，需要根据工作器的数量对数据集进行拆分，以确保模型收敛。输入的数据根据工作器其自身的索引来拆分，因此每个工作器各自负责处理该数据集 `1/num_workers` 个不同部分。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dma_wUAxZqo2"
      },
      "outputs": [],
      "source": [
        "BUFFER_SIZE = 10000\n",
        "BATCH_SIZE = 64\n",
        "\n",
        "def input_fn(mode, input_context=None):\n",
        "  datasets, info = tfds.load(name='mnist',\n",
        "                                with_info=True,\n",
        "                                as_supervised=True)\n",
        "  mnist_dataset = (datasets['train'] if mode == tf.estimator.ModeKeys.TRAIN else\n",
        "                   datasets['test'])\n",
        "\n",
        "  def scale(image, label):\n",
        "    image = tf.cast(image, tf.float32)\n",
        "    image /= 255\n",
        "    return image, label\n",
        "\n",
        "  if input_context:\n",
        "    mnist_dataset = mnist_dataset.shard(input_context.num_input_pipelines,\n",
        "                                        input_context.input_pipeline_id)\n",
        "  return mnist_dataset.map(scale).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4BlcVXMhB59T"
      },
      "source": [
        "使模型收敛的另一种合理方式是在每个工作器上设置不同的随机种子，然后对数据集进行随机重排。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8YFpxrcsZ2xG"
      },
      "source": [
        "## 多工作器配置\n",
        "\n",
        "本教程主要的不同（区别于[使用多 GPU 训练教程](./keras.ipynb)）在于多工作器的创建。明确集群中每个工作器的配置的标准方式是设置环境变量 `TF_CONFIG` 。\n",
        "\n",
        "`TF_CONFIG` 里包括了两个部分：`cluster` 和 `task`。`cluster` 提供了关于整个集群的信息，也就是集群中的工作器和参数服务器（parameter server）。`task` 提供了关于当前任务的信息。在本例中，任务的类型（type）是 worker 且该任务的索引（index）是 0。\n",
        "\n",
        "出于演示的目的，本教程展示了怎么将 `TF_CONFIG` 设置成两个本地的工作器。在实践中，你可以在外部的IP地址和端口上创建多个工作器，并为每个工作器正确地配置好 `TF_CONFIG` 变量，也就是更改任务的索引。\n",
        "\n",
        "警告：不要在 Colab 里执行以下代码。TensorFlow 的运行程序会试图在指定的 IP 地址和端口创建 gRPC 服务器，这会导致创建失败。\n",
        "\n",
        "```\n",
        "os.environ['TF_CONFIG'] = json.dumps({\n",
        "    'cluster': {\n",
        "        'worker': [\"localhost:12345\", \"localhost:23456\"]\n",
        "    },\n",
        "    'task': {'type': 'worker', 'index': 0}\n",
        "})\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qDreJzTffAP5"
      },
      "source": [
        "## 定义模型\n",
        "\n",
        "定义训练中用到的层，优化器和损失函数。本教程使用 Keras layers 定义模型，同[使用多 GPU 训练教程](./keras.ipynb)类似。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WNvOn_OeiUYC"
      },
      "outputs": [],
      "source": [
        "LEARNING_RATE = 1e-4\n",
        "def model_fn(features, labels, mode):\n",
        "  model = tf.keras.Sequential([\n",
        "      tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),\n",
        "      tf.keras.layers.MaxPooling2D(),\n",
        "      tf.keras.layers.Flatten(),\n",
        "      tf.keras.layers.Dense(64, activation='relu'),\n",
        "      tf.keras.layers.Dense(10, activation='softmax')\n",
        "  ])\n",
        "  logits = model(features, training=False)\n",
        "\n",
        "  if mode == tf.estimator.ModeKeys.PREDICT:\n",
        "    predictions = {'logits': logits}\n",
        "    return tf.estimator.EstimatorSpec(labels=labels, predictions=predictions)\n",
        "\n",
        "  optimizer = tf.compat.v1.train.GradientDescentOptimizer(\n",
        "      learning_rate=LEARNING_RATE)\n",
        "  loss = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "      from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(labels, logits)\n",
        "  loss = tf.reduce_sum(loss) * (1. / BATCH_SIZE)\n",
        "  if mode == tf.estimator.ModeKeys.EVAL:\n",
        "    return tf.estimator.EstimatorSpec(mode, loss=loss)\n",
        "\n",
        "  return tf.estimator.EstimatorSpec(\n",
        "      mode=mode,\n",
        "      loss=loss,\n",
        "      train_op=optimizer.minimize(\n",
        "          loss, tf.compat.v1.train.get_or_create_global_step()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P94PrIW_kSCE"
      },
      "source": [
        "注意：尽管在本例中学习率是固定的，但是通常情况下可能有必要基于全局的批次大小对学习率进行调整。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UhNtHfuxCGVy"
      },
      "source": [
        "## MultiWorkerMirroredStrategy\n",
        "\n",
        "为训练模型，需要使用 `tf.distribute.experimental.MultiWorkerMirroredStrategy` 实例。`MultiWorkerMirroredStrategy` 创建了每个设备中模型层里所有变量的拷贝，且是跨工作器的。其用到了 `CollectiveOps`，这是 TensorFlow 里的一种操作，用来整合梯度以及确保变量同步。该策略的更多细节可以在 [`tf.distribute.Strategy` 指南](../../guide/distribute_strategy.ipynb)中找到。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1uFSHCJXMrQ-"
      },
      "outputs": [],
      "source": [
        "strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H47DDcOgfzm7"
      },
      "source": [
        "## 训练和评估模型\n",
        "接下来，在 `RunConfig` 中为 estimator 指明分布式策略，同时通过调用 `tf.estimator.train_and_evaluate` 训练和评估模型。本教程只通过指明 `train_distribute` 进行分布式训练。但是也同样也可以通过指明 `eval_distribute` 来进行分布式评估。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BcsuBYrpgnlS"
      },
      "outputs": [],
      "source": [
        "config = tf.estimator.RunConfig(train_distribute=strategy)\n",
        "\n",
        "classifier = tf.estimator.Estimator(\n",
        "    model_fn=model_fn, model_dir='/tmp/multiworker', config=config)\n",
        "tf.estimator.train_and_evaluate(\n",
        "    classifier,\n",
        "    train_spec=tf.estimator.TrainSpec(input_fn=input_fn),\n",
        "    eval_spec=tf.estimator.EvalSpec(input_fn=input_fn)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XVk4ftYx6JAO"
      },
      "source": [
        "# 优化训练后的模型性能\n",
        "\n",
        "现在你已经有了由 `tf.distribute.Strategy` 的模型和能支持多工作器的 Estimator。你可以尝试使用下列技巧来优化多工作器训练的性能。\n",
        "\n",
        "*   *增加单批次的大小：* 此处的批次大小指的是每个 GPU 上的批次大小。通常来说，最大的批次大小应该适应 GPU 的内存大小。\n",
        "*   *变量转换：* 尽可能将变量转换成 `tf.float`。官方的 ResNet 模型包括了如何完成的[样例](https://github.com/tensorflow/models/blob/8367cf6dabe11adf7628541706b660821f397dce/official/resnet/resnet_model.py#L466)。\n",
        "*   *使用集群通信：* `MultiWorkerMirroredStrategy` 提供了好几种[集群通信的实现](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/distribute/cross_device_ops.py).  \n",
        "    * `RING` 实现了基于环状的集群，使用了 gRPC 作为跨主机通讯层。  \n",
        "    * `NCCL` 使用了 [英伟达的 NCCL](https://developer.nvidia.com/nccl) 来实现集群。\n",
        "    * `AUTO` 将选择延后至运行时。\n",
        "\n",
        "集群实现的最优选择不仅基于 GPU 的数量和种类，也基于集群间的通信网络。想要覆盖自动的选项，需要指明 `MultiWorkerMirroredStrategy` 的构造器里的 `communication` 参数，例如让 `communication=tf.distribute.experimental.CollectiveCommunication.NCCL` 。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AW0Hb2xM6EGX"
      },
      "source": [
        "## 更多的代码示例\n",
        "\n",
        "1. [端到端的示例](https://github.com/tensorflow/ecosystem/tree/master/distribution_strategy)里使用了 Kubernetes 模板。在这个例子里我们一开始使用了 Keras 模型，并使用了 `tf.keras.estimator.model_to_estimator` API 将其转换成了 Estimator。\n",
        "2. 官方的 [ResNet50](https://github.com/tensorflow/models/blob/master/official/resnet/imagenet_main.py) 模型，我们可以使用 `MirroredStrategy` 或 `MultiWorkerMirroredStrategy` 来训练它。"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "Tce3stUlHN0L"
      ],
      "name": "multi_worker_with_estimator.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
